Amazon SageMaker Debuggerによって収集したデータをTensorBoardで確認する

osawa yuto

2020.01.30

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

TensorBoardはTensorFlowが提供する可視化ツールで、データソースとしてS3上のデータにも対応しています。今回はSageMaker Debuggerを使って収集したテンソルデータをTensorBoard用にS3に出力し、その内容をTensorBoardから確認してみたいと思います。

概要

TensorBoardOutputConfigを次のように設定することで、SageMaker Debuggerが収集したテンソルデータをTensorBoardで参照できるように自動で加工し、保存してくれます。

from sagemaker.tensorflow import TensorFlow
from sagemaker.debugger import DebuggerHookConfig, TensorBoardOutputConfig, CollectionConfig
estimator = TensorFlow(
    entry_point='./path/to/script', #  学習用スクリプト
    tensorboard_output_config=TensorBoardOutputConfig('s3://bucket/path/to/log'), # TensorBoardの参照用のデータ保存場所
    debugger_hook_config=DebuggerHookConfig(
        s3_output_path='s3://bucket/path/to/tensor, # 収集したテンソルの保存場所
        collection_configs=[ # 収集したいテンソル
            CollectionConfig(name="weights"),
            CollectionConfig(name="gradients"),
            CollectionConfig(name="losses"),
            ...
        ]
    ),
    ...
)

TensorBoard用にデータが保存されたパスをTensorBoardで指定することで、内容を確認出来ます。

export AWS_ACCESS_KEY_ID=HOGEHOGEHOGE
export AWS_SECRET_ACCESS_KEY=aaaaaaaaaaaa/aaaaaa
# export AWS_SESSION_TOKEN=AAAAAAAAAA.....  # 必要に応じてAWS_SESSION_TOKENも設定
export AWS_REGION=ap-northeast-1 # 対象のS3バケットがあるリージョン
tensorboard --logdir s3://bucket/path/to/log

対応バージョン

SageMaker DebuggerはTensorFlow 1.13, 1.14, 1.15とKeras 2.3に対応しています。(2020年1月30日現在)

sagemaker-debugger/tensorflow.md at master · awslabs/sagemaker-debugger

ゼロスクリプトチェンジ

以下のフレームワークバージョンのコンテナイメージを利用することで、エントリポイントとして設定するスクリプトを変更することなく、SageMaker Debuggerを使用できます。

TensorFlow1.15
- Prebuilt Amazon SageMaker Docker Images for TensorFlow, MXNet, Chainer, and PyTorch - Amazon SageMaker
- Docker Images for Deep Learning - AWS

やってみる

Amazon SageMaker Examplesで紹介されているノートブックの内容を元に試してみます。

学習

まずは、使用するライブラリ/モジュールを読み込み、学習に関連する情報を定義しておきます。

import boto3
from datetime import datetime
from os import path
import sagemaker
from sagemaker.tensorflow import TensorFlow
from sagemaker.debugger import DebuggerHookConfig, TensorBoardOutputConfig, CollectionConfig

# 保存場所
bucket = 'bucket_name'
s3_prefix = f's3://{bucket}/tensorflow-job-test-' + datetime.now().strftime('%Y%m%d-%H%M%S')

# 学習用スクリプト
entrypoint_script='src/mnist_zerocodechange.py'

今回はエントリポイントとして設定する学習用スクリプトは以下のものを使用します。MNISTのデータセットを読み込み、CNNの分類モデルを学習します。このスクリプトには、SageMaker Debugger用にデータを収集するためにフックを設定するといった処理はありません。今回はSageMaker Debuggerにネイティブで対応しているコンテナイメージを使用するため、専用の処理を追加することなく、学習中のデータを収集し、デバッグすることが可能です。

学習用スクリプト: amazon-sagemaker-examples/mnist_zerocodechange.py at master · awslabs/amazon-sagemaker-examples

次に、収集するテンソルデータの種類や保存場所とTensorBoard用のデータの保存場所を設定します。

log_uri = path.join(s3_prefix, 'log')
tb_config = TensorBoardOutputConfig(log_uri)
hook_config = DebuggerHookConfig(
    s3_output_path=s3_prefix,
    collection_configs=[
        CollectionConfig(name="weights"),
        CollectionConfig(name="gradients"),
        CollectionConfig(name="losses")
    ]
)

TensorFlow用のEstimatorに学習に関する各種パラメータを設定し、学習ジョブを作成します。

estimator = TensorFlow(
    entry_point=entrypoint_script,
    tensorboard_output_config=tb_config,
    debugger_hook_config=hook_config
    role=sagemaker.get_execution_role(),
    base_job_name='smdebugger-demo-mnist-tensorflow',
    train_instance_count=1,
    train_instance_type='ml.m4.xlarge',
    train_volume_size=400,
    framework_version='1.15',
    py_version='py3',
    train_max_run=3600,
    script_mode=True,
    hyperparameters={"num_epochs": 3},
    output_path=s3_prefix
)
estimator.fit()

TensorBoard

認証情報等を設定し、TensorBoardを立ち上げます。

export AWS_ACCESS_KEY_ID=HOGEHOGEHOGE
export AWS_SECRET_ACCESS_KEY=aaaaaaaaaaaa/aaaaaa
# export AWS_SESSION_TOKEN=AAAAAAAAAA.....  # 必要に応じてAWS_SESSION_TOKENも設定
export AWS_REGION=ap-northeast-1 # 対象のS3バケットがあるリージョン
tensorboard --logdir {log_uri}

http://localhost:6006/ でTensorBoardを開きます。ポート番号は環境によっては変わる可能性があります。

トップのSCALARSではロスの変化が表示されています。

GRAPHSからモデルの構成を確認できます。

DISTRIBUTIONSで収集したデータの分布を見れます。

HISTOGRAMSで収集したデータのヒストグラムを確認できます。

さいごに

SageMaker Debuggerを使って収集したテンソルデータをTensorBoardで確認する方法について紹介しました。学習用のスクリプトを変更せずに、モデルの学習状況を収集し、分析することができるのはかなり便利ではないでしょうか。TensorFlow 2.0や2.1といったより新しいバージョンへのSageMaker Debuggerの対応も今後に期待したいところです。